import cv2
import numpy as np

# ---------------------- 可调参数（全局） ----------------------
VIDEO_PATH = '动画表情视频.mp4'                                  # 视频路径
TEMPLATE_PATHS = [                                               # 模板路径列表
    'small_template_1.png',
    'small_template_2.png',
    'small_template_3.png'
]
DISPLAY_WIDTH = 540                                              # 显示窗口宽度
DISPLAY_HEIGHT = 960                                             # 显示窗口高度
MATCH_STEP_X = 40                                                # 横向滑窗步长
MATCH_STEP_Y = 40                                                # 纵向滑窗步长
MATCH_STEP_X_FINE = 20                                           # 失踪恢复时的细粒度滑窗（横）
MATCH_STEP_Y_FINE = 20                                           # 失踪恢复时的细粒度滑窗（纵）
LOCAL_SEARCH_RADIUS = 10                                         # 局部搜索半径（像素）
SCORE_THRESHOLD = 0.34                                           # 匹配得分阈值
SCALES = [0.3,0.5,0.7, 0.9, 1.1, 1.3]                            # 缩放比例列表
FRAME_DOWNSAMPLE_RATIO = 0.5                                     # 图像尺寸缩放比例（加速匹配）
SKIP_FRAMES = 0                                                  # 每隔多少帧进行一次全局搜索（0 表示每帧都做）
RETRY_GLOBAL_ON_FAIL = True                                      # 失败后是否允许自动回退到全局搜索

# ---------------------- NCC计算函数 ----------------------
def compute_ncc(patch, template):
    """
    计算归一化互相关（Normalized Cross-Correlation）得分
    输入：
      patch: 当前图像块（numpy数组）
      template: 模板图像（numpy数组）
    输出：
      NCC得分（范围[-1,1]），值越大相似度越高
    """
    # 类型转换与数据标准化
    patch = patch.astype(np.float32)
    template = template.astype(np.float32)
    # 形状校验
    if patch.shape != template.shape:
        return -1.0
    # 均值中心化
    patch_mean = np.mean(patch)
    template_mean = np.mean(template)
    # 分子：协方差计算
    numerator = np.sum((patch - patch_mean) * (template - template_mean))
    # 分母：标准差乘积
    denominator = np.sqrt(np.sum((patch - patch_mean) ** 2) *np.sum((template - template_mean) ** 2))
    # 防止除以零
    if denominator == 0:
        return -1.0
    return numerator / denominator  # 归一化得分

# ---------------------- 局部搜索 ----------------------
def local_search(frame_gray, center, template, search_radius=LOCAL_SEARCH_RADIUS):
    """
    基于滑动窗口的局部最优匹配搜索
    输入：
      frame_gray: 当前帧灰度图像
      center: 上一帧目标中心坐标
      template: 当前模板
      search_radius: 搜索半径（像素）
    输出：
      best_location: 最佳匹配位置坐标 (x,y)
      best_score: 最高匹配得分
    """
    h, w = template.shape
    cx, cy = center
    best_score = -1.0
    best_location = None
    # 双循环遍历局部搜索区域
    for dy in range(-search_radius, search_radius + 1):
        for dx in range(-search_radius, search_radius + 1):
            # 候选位置计算
            x = int(cx - w // 2 + dx)
            y = int(cy - h // 2 + dy)
            # 边界检查
            if y < 0 or x < 0 or y + h > frame_gray.shape[0] or x + w > frame_gray.shape[1]:
                continue
            # 提取候选区域并计算相似度
            patch = frame_gray[y:y + h, x:x + w]
            score = compute_ncc(patch, template)
            # 更新最优结果
            if score > best_score:
                best_score = score
                best_location = (x, y)
    return best_location, best_score

# ---------------------- 生成缩放模板 ----------------------
def generate_scaled_templates(template):
    """
    生成不同缩放比例的模板集合
    输入：
      template: 原始模板图像
    输出：
      scaled_templates: 包含(scale_factor, scaled_template)的元组列表
    """
    scaled_templates = []
    for scale in SCALES:
        # 计算新尺寸
        h, w = template.shape
        new_w, new_h = int(w * scale), int(h * scale)
        # 双线性插值缩放
        scaled = cv2.resize(template, (new_w, new_h),
                            interpolation=cv2.INTER_LINEAR)
        scaled_templates.append((scale, scaled))
    return scaled_templates

# ---------------------- 主函数 ----------------------
def main():
    # 读取视频帧
    cap = cv2.VideoCapture(VIDEO_PATH)
    # 灰度化与下采样（加速计算）
    templates = [cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2GRAY) for path in TEMPLATE_PATHS]
    templates = [cv2.resize(t, (0, 0), fx=FRAME_DOWNSAMPLE_RATIO, fy=FRAME_DOWNSAMPLE_RATIO) for t in templates]

    tracking = False
    track_center = None
    trajectory = []
    current_template_idx = 0
    current_scale = 1.0
    template = templates[current_template_idx]
    frame_count = 0
    force_global_search = False

    #设置可视化视频大小
    cv2.namedWindow('Tracking', cv2.WINDOW_NORMAL)
    cv2.resizeWindow('Tracking', DISPLAY_WIDTH, DISPLAY_HEIGHT)

    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame_count += 1

        frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        frame_gray = cv2.resize(frame_gray, (0, 0), fx=FRAME_DOWNSAMPLE_RATIO, fy=FRAME_DOWNSAMPLE_RATIO)

        need_global_search = not tracking or force_global_search or (SKIP_FRAMES > 0 and frame_count % SKIP_FRAMES == 0)
        force_global_search = False

        # 判断是否需要进行全局搜索
        if need_global_search:
            best_score = -1.0
            best_loc = None
            best_template = None
            best_scale = 1.0
            best_idx = -1

            step_x = MATCH_STEP_X_FINE if not tracking else MATCH_STEP_X
            step_y = MATCH_STEP_Y_FINE if not tracking else MATCH_STEP_Y
            # 多模板、多尺度全局搜索
            for idx, base_template in enumerate(templates):
                for scale, scaled_template in generate_scaled_templates(base_template):
                    # 进行缩放操作，生成缩放模板
                    th, tw = scaled_template.shape

                    # 滑动窗口遍历
                    for y in range(0, frame_gray.shape[0] - th + 1, step_y):
                        for x in range(0, frame_gray.shape[1] - tw + 1, step_x):
                            # 计算NCC得分
                            patch = frame_gray[y:y+th, x:x+tw]
                            score = compute_ncc(patch, scaled_template)
                            # 记录最佳匹配
                            if score > best_score:
                                best_score = score
                                best_loc = (x, y)
                                best_template = scaled_template
                                best_scale = scale
                                best_idx = idx

            if best_score > SCORE_THRESHOLD:
                x, y = best_loc
                template = best_template
                current_template_idx = best_idx
                current_scale = best_scale
                track_center = (x + template.shape[1] // 2, y + template.shape[0] // 2)
                tracking = True
                print(f"[INIT MATCH] Frame {frame_count}, Template #{current_template_idx + 1}, Scale={current_scale:.2f}, NCC Score={best_score:.3f}")

        else:
            # 局部优化搜索
            loc, score = local_search(frame_gray, track_center, template)
            print(f"[LOCAL SEARCH] Frame {frame_count}, Template #{current_template_idx + 1}, Scale={current_scale:.2f}, NCC Score={score:.3f}")
            if loc and score > SCORE_THRESHOLD:
                x, y = loc
                th, tw = template.shape
                track_center = (x + tw // 2, y + th // 2)
            else:
                improved = False
                best_scaled_score = -1.0
                best_scaled_loc = None
                best_scaled_template = None
                for scale, scaled_template in generate_scaled_templates(template):
                    loc, scaled_score = local_search(frame_gray, track_center, scaled_template)
                    if loc and scaled_score > best_scaled_score:
                        x, y = loc
                        center_candidate = (x + scaled_template.shape[1] // 2, y + scaled_template.shape[0] // 2)
                        dist = np.linalg.norm(np.array(center_candidate) - np.array(track_center))
                        if scaled_score > SCORE_THRESHOLD and dist < 30:
                            best_scaled_score = scaled_score
                            best_scaled_loc = loc
                            best_scaled_template = scaled_template
                            current_scale = scale
                            improved = True
                # 失败处理策略
                if improved:
                    template = best_scaled_template
                    x, y = best_scaled_loc
                    th, tw = template.shape
                    track_center = (x + tw // 2, y + th // 2)
                else:
                    tracking = False
                    if RETRY_GLOBAL_ON_FAIL:
                        force_global_search = True

        # 轨迹绘制与可视化
        if tracking:
            trajectory.append(track_center)
            x, y = int(track_center[0] - template.shape[1] // 2), int(track_center[1] - template.shape[0] // 2)
            frame_display = cv2.resize(frame, (0, 0), fx=FRAME_DOWNSAMPLE_RATIO, fy=FRAME_DOWNSAMPLE_RATIO)
            cv2.rectangle(frame_display, (x, y), (x + template.shape[1], y + template.shape[0]), (0, 0, 255), 2)
        else:
            frame_display = cv2.resize(frame, (0, 0), fx=FRAME_DOWNSAMPLE_RATIO, fy=FRAME_DOWNSAMPLE_RATIO)

        for i in range(1, len(trajectory)):
            cv2.line(frame_display, trajectory[i - 1], trajectory[i], (0, 0, 255), 2)

        info_text = f"Template: {current_template_idx + 1}"
        cv2.putText(frame_display, info_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2)

        frame_resized = cv2.resize(frame_display, (DISPLAY_WIDTH, DISPLAY_HEIGHT))
        cv2.imshow('Tracking', frame_resized)
        if cv2.waitKey(30) & 0xFF == 27:
            break

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()
